from prune_utils.prune_utils import *
import torch


def ObtainBnMask(bn_module, threshold):
    thre = threshold.cuda()
    mask = bn_module.weight.data.abs().ge(thre).float()
    return mask


def GetPruneThre(model, prune_idx, percent=0.8):
    """
    获取裁剪阈值，大概思路，先对bn的scale进行从小到大的排序，然后根据裁剪比例获取对应index的阈值
    :param model:
    :param percent:
    :return:
    """
    model_copy=deepcopy(model)
    bn_weights = GatherBnWeights(model_copy.models, prune_idx)
    sorted_bn = torch.sort(bn_weights)[0]
    thre_index = int(len(sorted_bn) * percent)
    thre_value = sorted_bn[thre_index]
    print(f"增强比例值小于{thre_value:.6f} 将被抑制！")
    # 计算当前的实际压缩比例值,其中真正的压缩比例值是实际的小的，但是小多少，这个是未知的，因为在这个index未知的
    # 值可能有很多个
    remain_nums = 0
    for idx in prune_idx:
        bn_module = model_copy.models[idx][1]
        mask = ObtainBnMask(bn_module, thre_value)
        remain_nums += int(mask.sum())
        bn_module.weight.data.mul_(mask)
    # 此处可以进行测试下现模型的基本准确度，一般会比较低，如果降的很低，证明这个
    # 阈值剪掉了一些还有意义的通道
    print(f'Number of channels has been reduced from {len(sorted_bn)} to {remain_nums}')
    print(f'Prune ratio:{1 - remain_nums / len(sorted_bn):.3f}')
    return thre_value


def GetFiltersMask(model, thre_value, CBL_idx, prune_idx, hold=0.3):
    '''
    获取压缩后的网络配置
    :param model:
    :param thre_value:
    :param CBL_idx:
    :param prune_idx:
    :param hold:
    :return:
    '''
    pruned = 0
    total = 0
    num_filters = []
    filters_mask = []
    for idx in CBL_idx:
        bn_module = model.models[idx][1]
        if idx in prune_idx:
            mask = ObtainBnMask(bn_module, thre_value).cpu().numpy()
            remain = int(mask.sum())
            # pruned=pruned
            if remain == 0:
                # max_value=bn_module.weight.data.abs().max()
                # 保留原始值的hold比例内的值
                sort_w = torch.sort(bn_module.weight.data.abs())
                sort_index = len(sort_w) * hold
                new_thre_value = sort_w[sort_index]
                mask = ObtainBnMask(bn_module, new_thre_value).cpu().numpy()
                remain = int(mask.sum())

            pruned = pruned + mask.shape[0] - remain
            print(f'layer index:{idx:>3d} \t total channel:{mask.shape[0]:>4d}\t'
                  f'remaining channel:{remain:>4d}')
        else:
            mask = np.ones(bn_module.weight.data.shape)
            remain = mask.shape[0]
        total += mask.shape[0]
        num_filters.append(remain)
        filters_mask.append(mask.copy())

    prune_ratio = pruned / total
    print(f'Prune channels:{pruned}\t Prune ratio:{prune_ratio:.3f}')
    return num_filters, filters_mask
